import os
import json
import logging
from typing import List, Dict, Any, Optional
from datetime import datetime
from elasticsearch import Elasticsearch
from elasticsearch.helpers import bulk

logger = logging.getLogger(__name__)

class ElasticsearchService:
    """Elasticsearch服务类 - 兼容ES 7.17.18版本"""
    
    def __init__(self):
        # ES连接配置 - 使用阿里云ES地址
        self.es_host = os.getenv("ELASTICSEARCH_HOST", "114.55.65.49")
        self.es_port = int(os.getenv("ELASTICSEARCH_PORT", "9200"))
        self.es_username = os.getenv("ELASTICSEARCH_USERNAME", "")
        self.es_password = os.getenv("ELASTICSEARCH_PASSWORD", "")
        
        # 索引配置
        self.index_name = "didi_knowledge_base"
        
        # 根据ES 7.17.18版本调整索引设置
        self.index_settings = self._get_index_settings()
        
        # 初始化ES客户端
        self.es = self._init_elasticsearch()
        
    def _get_index_settings(self):
        """根据ES 7.17.18版本获取索引设置"""
        # ES 7.17.18兼容的设置
        settings = {
            "settings": {
                "number_of_shards": 1,
                "number_of_replicas": 0,
                "analysis": {
                    "analyzer": {
                        "chinese_analyzer": {
                            "type": "custom",
                            "tokenizer": "standard",
                            "filter": ["lowercase", "stop"]
                        }
                    }
                }
            },
            "mappings": {
                "properties": {
                    "title": {
                        "type": "text",
                        "analyzer": "chinese_analyzer",
                        "search_analyzer": "chinese_analyzer"
                    },
                    "content": {
                        "type": "text",
                        "analyzer": "chinese_analyzer",
                        "search_analyzer": "chinese_analyzer"
                    },
                    "category": {
                        "type": "keyword"
                    },
                    "url": {
                        "type": "keyword"
                    },
                    "tags": {
                        "type": "keyword"
                    },
                    "priority": {
                        "type": "integer"
                    },
                    "crawl_time": {
                        "type": "date"
                    },
                    "processed_time": {
                        "type": "date"
                    },
                    "word_count": {
                        "type": "integer"
                    },
                    "source": {
                        "type": "keyword"
                    },
                    "chunk_id": {
                        "type": "integer"
                    }
                }
            }
        }
        
        return settings
        
    def _init_elasticsearch(self) -> Elasticsearch:
        """初始化Elasticsearch客户端 - ES 7.17.18兼容"""
        try:
            # 构建ES连接URL
            es_url = f"http://{self.es_host}:{self.es_port}"
            logger.info(f"尝试连接Elasticsearch: {es_url}")
            
            # ES 7.17.18的连接方式 - 使用配置的地址
            if self.es_username and self.es_password:
                es = Elasticsearch(
                    [es_url],
                    http_auth=(self.es_username, self.es_password),
                    request_timeout=30,
                    max_retries=3,
                    retry_on_timeout=True
                )
            else:
                es = Elasticsearch(
                    [es_url],
                    request_timeout=30,
                    max_retries=3,
                    retry_on_timeout=True
                )
            
            # 测试连接
            if es.ping():
                logger.info("Elasticsearch连接成功")
                # 获取版本信息
                info = es.info()
                version = info.get('version', {}).get('number', 'unknown')
                logger.info(f"Elasticsearch版本: {version}")
                return es
            else:
                logger.error("Elasticsearch连接失败")
                return None
                
        except Exception as e:
            logger.error(f"Elasticsearch初始化失败: {e}")
            return None
    
    def create_index(self) -> bool:
        """创建索引 - ES 7.17.18兼容"""
        try:
            if not self.es:
                logger.error("Elasticsearch客户端未初始化")
                return False
            
            # 检查索引是否存在
            if self.es.indices.exists(index=self.index_name):
                logger.info(f"索引 {self.index_name} 已存在")
                return True
            
            # 创建索引
            response = self.es.indices.create(
                index=self.index_name,
                body=self.index_settings
            )
            
            if response.get('acknowledged', False):
                logger.info(f"索引 {self.index_name} 创建成功")
                return True
            else:
                logger.error(f"索引 {self.index_name} 创建失败")
                return False
                
        except Exception as e:
            logger.error(f"创建索引失败: {e}")
            return False
    
    def index_documents(self, documents: List[Dict[str, Any]]) -> bool:
        """索引文档 - ES 7.17.18兼容"""
        try:
            if not self.es:
                logger.error("Elasticsearch客户端未初始化")
                return False
            
            # 确保索引存在
            if not self.create_index():
                return False
            
            # 准备批量索引数据
            actions = []
            for doc in documents:
                action = {
                    "_index": self.index_name,
                    "_source": {
                        "title": doc.get("title", ""),
                        "content": doc.get("content", ""),
                        "category": doc.get("category", ""),
                        "url": doc.get("url", ""),
                        "tags": doc.get("tags", []),
                        "priority": doc.get("priority", 1),
                        "crawl_time": doc.get("crawl_time", datetime.now().isoformat()),
                        "processed_time": doc.get("processed_time", datetime.now().isoformat()),
                        "word_count": doc.get("word_count", 0),
                        "source": doc.get("source", ""),
                        "chunk_id": doc.get("chunk_id", 0)
                    }
                }
                actions.append(action)
            
            # 批量索引 - ES 7.17.18兼容
            success, failed = bulk(self.es, actions, chunk_size=1000, request_timeout=30)
            
            logger.info(f"索引完成: 成功 {success} 个文档, 失败 {len(failed)} 个文档")
            return len(failed) == 0
            
        except Exception as e:
            logger.error(f"索引文档失败: {e}")
            return False
    
    def search_documents(self, query: str, category: str = None, 
                        size: int = 10, from_: int = 0) -> List[Dict[str, Any]]:
        """搜索文档 - ES 7.17.18兼容"""
        try:
            if not self.es:
                logger.error("Elasticsearch客户端未初始化")
                return []
            
            # 构建搜索查询
            search_body = {
                "query": {
                    "bool": {
                        "must": [
                            {
                                "multi_match": {
                                    "query": query,
                                    "fields": ["title^2", "content"],
                                    "type": "best_fields",
                                    "fuzziness": "AUTO"
                                }
                            }
                        ]
                    }
                },
                "highlight": {
                    "fields": {
                        "title": {},
                        "content": {
                            "fragment_size": 150,
                            "number_of_fragments": 3
                        }
                    }
                },
                "sort": [
                    {"_score": {"order": "desc"}},
                    {"priority": {"order": "desc"}},
                    {"crawl_time": {"order": "desc"}}
                ],
                "size": size,
                "from": from_
            }
            
            # 如果指定了类别，添加过滤条件
            if category:
                search_body["query"]["bool"]["filter"] = [
                    {"term": {"category": category}}
                ]
            
            # 执行搜索
            response = self.es.search(
                index=self.index_name,
                body=search_body
            )
            
            # 处理搜索结果
            results = []
            for hit in response['hits']['hits']:
                result = {
                    "id": hit['_id'],
                    "score": hit['_score'],
                    "title": hit['_source'].get('title', ''),
                    "content": hit['_source'].get('content', ''),
                    "category": hit['_source'].get('category', ''),
                    "url": hit['_source'].get('url', ''),
                    "tags": hit['_source'].get('tags', []),
                    "priority": hit['_source'].get('priority', 1),
                    "source": hit['_source'].get('source', ''),
                    "highlight": hit.get('highlight', {})
                }
                results.append(result)
            
            return results
            
        except Exception as e:
            logger.error(f"搜索文档失败: {e}")
            return []
    
    def search_by_category(self, category: str, size: int = 20) -> List[Dict[str, Any]]:
        """按类别搜索文档 - ES 7.17.18兼容"""
        try:
            if not self.es:
                return []
            
            search_body = {
                "query": {
                    "term": {
                        "category": category
                    }
                },
                "sort": [
                    {"priority": {"order": "desc"}},
                    {"crawl_time": {"order": "desc"}}
                ],
                "size": size
            }
            
            response = self.es.search(
                index=self.index_name,
                body=search_body
            )
            
            results = []
            for hit in response['hits']['hits']:
                result = {
                    "id": hit['_id'],
                    "title": hit['_source'].get('title', ''),
                    "content": hit['_source'].get('content', ''),
                    "category": hit['_source'].get('category', ''),
                    "url": hit['_source'].get('url', ''),
                    "priority": hit['_source'].get('priority', 1),
                    "source": hit['_source'].get('source', '')
                }
                results.append(result)
            
            return results
            
        except Exception as e:
            logger.error(f"按类别搜索失败: {e}")
            return []
    
    def get_statistics(self) -> Dict[str, Any]:
        """获取索引统计信息 - ES 7.17.18兼容"""
        try:
            if not self.es:
                return {"error": "Elasticsearch客户端未初始化"}
            
            # 获取索引统计
            stats = self.es.indices.stats(index=self.index_name)
            
            # 获取文档数量
            count_response = self.es.count(index=self.index_name)
            doc_count = count_response['count']
            
            # 按类别统计
            aggs_body = {
                "size": 0,
                "aggs": {
                    "categories": {
                        "terms": {
                            "field": "category",
                            "size": 20
                        }
                    },
                    "sources": {
                        "terms": {
                            "field": "source",
                            "size": 10
                        }
                    }
                }
            }
            
            aggs_response = self.es.search(
                index=self.index_name,
                body=aggs_body
            )
            
            categories = {}
            for bucket in aggs_response['aggregations']['categories']['buckets']:
                categories[bucket['key']] = bucket['doc_count']
            
            sources = {}
            for bucket in aggs_response['aggregations']['sources']['buckets']:
                sources[bucket['key']] = bucket['doc_count']
            
            return {
                "total_documents": doc_count,
                "categories": categories,
                "sources": sources,
                "index_size_bytes": stats['indices'][self.index_name]['total']['store']['size_in_bytes'],
                "last_update": datetime.now().isoformat()
            }
            
        except Exception as e:
            logger.error(f"获取统计信息失败: {e}")
            return {"error": str(e)}
    
    def delete_document(self, doc_id: str) -> bool:
        """删除文档 - ES 7.17.18兼容"""
        try:
            if not self.es:
                return False
            
            response = self.es.delete(
                index=self.index_name,
                id=doc_id
            )
            
            return response['result'] == 'deleted'
            
        except Exception as e:
            logger.error(f"删除文档失败: {e}")
            return False
    
    def update_document(self, doc_id: str, update_data: Dict[str, Any]) -> bool:
        """更新文档 - ES 7.17.18兼容"""
        try:
            if not self.es:
                return False
            
            response = self.es.update(
                index=self.index_name,
                id=doc_id,
                body={"doc": update_data}
            )
            
            return response['result'] == 'updated'
            
        except Exception as e:
            logger.error(f"更新文档失败: {e}")
            return False
    
    def suggest_search(self, query: str, size: int = 5) -> List[str]:
        """搜索建议 - ES 7.17.18兼容"""
        try:
            if not self.es:
                return []
            
            suggest_body = {
                "suggest": {
                    "title_suggest": {
                        "prefix": query,
                        "completion": {
                            "field": "title_suggest",
                            "size": size
                        }
                    }
                }
            }
            
            response = self.es.search(
                index=self.index_name,
                body=suggest_body
            )
            
            suggestions = []
            for suggestion in response['suggest']['title_suggest']:
                for option in suggestion['options']:
                    suggestions.append(option['text'])
            
            return suggestions
            
        except Exception as e:
            logger.error(f"搜索建议失败: {e}")
            return []
    
    def reindex_data(self, documents: List[Dict[str, Any]]) -> bool:
        """重新索引数据 - ES 7.17.18兼容"""
        try:
            # 删除现有索引
            if self.es.indices.exists(index=self.index_name):
                self.es.indices.delete(index=self.index_name)
                logger.info(f"删除索引 {self.index_name}")
            
            # 重新创建索引
            if not self.create_index():
                return False
            
            # 重新索引文档
            return self.index_documents(documents)
            
        except Exception as e:
            logger.error(f"重新索引失败: {e}")
            return False

 